package main

import (
	"fmt"
	"sync"
)

// 参考文章: https://segmentfault.com/a/1190000006261218
func gen(done <-chan struct{}, nums ...int) <-chan int {
	out := make(chan int)
	go func() {
		for _, n := range nums {
			select {
			case out <- n:
			case <-done:

			}
		}
		defer close(out)
	}()
	return out
}

func sq(done <-chan struct{}, in <-chan int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for n := range in {
			select {
			case out <- n * n:
			case <-done:
				return
			}
		}
	}()
	return out
}

func merge(done <-chan struct{}, cs ...<-chan int) <-chan int {
	var wg sync.WaitGroup
	out := make(chan int)

	// 为 cs 的每一个 channel 创建一个 goroutine
	// 这个 goroutine 运行 output，它将数据从 c
	// 拷贝到 out，直到 c 关闭，或者 接收到 done
	// 的关闭信号。人啊后调用 wg.Done()
	output := func(c <-chan int) {
		for n := range c {
			select {
			case out <- n:
			case <-done:
				return
			}
		}
		defer wg.Done()
	}
	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

	// 启动一个 goroutine，用于所有 output goroutine结束时，关闭 out
	// 该goroutine 必须在 wg.Add 之后启动
	go func() {
		wg.Wait()
		close(out)
	}()

	return out

}

func main() {
	// 设置一个 全局共享的 done channel，
	// 当流水线退出时，关闭 done channel
	// 所有 goroutine接收到 done 的信号后，
	// 都会正常退出。
	done := make(chan struct{})
	defer close(done)

	in := gen(done, 2, 3, 100, 2, 3, 4, 6)
	// 将 sq 的工作分发给两个goroutine
	// 这两个 goroutine 均从 in 读取数据
	c1 := sq(done, in)
	c2 := sq(done, in)

	// 消费 outtput 生产的第一个值
	out := merge(done, c1, c2)
	fmt.Println(<-out)
	fmt.Println(<-out)
	fmt.Println(<-out)
	fmt.Println(<-out)

	// defer 调用时，done channel 会被关闭。
}
